{
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "version": "3.6.4",
      "file_extension": ".py",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "name": "python",
      "mimetype": "text/x-python"
    },
    "kaggle": {
      "accelerator": "nvidiaTeslaT4",
      "dataSources": [],
      "dockerImageVersionId": 30840,
      "isInternetEnabled": true,
      "language": "python",
      "sourceType": "notebook",
      "isGpuEnabled": true
    },
    "colab": {
      "name": "Titans Paper Implantation",
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "nbformat_minor": 0,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/codewithdark-git/Titans_Paper_Implementation/blob/main/Titans_Paper_Implantation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torch torchvision torchaudio"
      ],
      "metadata": {
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gcfnivgCn7xW",
        "outputId": "525f0ba4-33d3-4bce-aab8-a323a105e8e5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: torch in /usr/local/lib/python3.11/dist-packages (2.5.1+cu121)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.11/dist-packages (0.20.1+cu121)\n",
            "Requirement already satisfied: torchaudio in /usr/local/lib/python3.11/dist-packages (2.5.1+cu121)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch) (3.16.1)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.11/dist-packages (from torch) (4.12.2)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch) (3.4.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch) (3.1.5)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch) (2024.9.0)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.11/dist-packages (from torch) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.11/dist-packages (from torch) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.11/dist-packages (from torch) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.11/dist-packages (from torch) (9.1.0.70)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.11/dist-packages (from torch) (12.1.3.1)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.11/dist-packages (from torch) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.11/dist-packages (from torch) (10.3.2.106)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.11/dist-packages (from torch) (11.4.5.107)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.11/dist-packages (from torch) (12.1.0.106)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from torch) (2.21.5)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.11/dist-packages (from torch) (12.1.105)\n",
            "Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.11/dist-packages (from torch) (3.1.0)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.11/dist-packages (from torch) (1.13.1)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.11/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch) (12.6.85)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy==1.13.1->torch) (1.3.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.11/dist-packages (from torchvision) (1.26.4)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.11/dist-packages (from torchvision) (11.1.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch) (3.0.2)\n"
          ]
        }
      ],
      "execution_count": 1
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import math\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "\n",
        "class TitansMemoryModule(nn.Module):\n",
        "    def __init__(self, d_model, memory_size=512):\n",
        "        super().__init__()\n",
        "        self.memory_size = memory_size\n",
        "        self.memory = nn.Parameter(torch.zeros(memory_size, d_model))\n",
        "        self.key_proj = nn.Linear(d_model, d_model)\n",
        "        self.value_proj = nn.Linear(d_model, d_model)\n",
        "        self.forgetting_gate = nn.Linear(d_model, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # x shape: [batch_size, seq_len, d_model]\n",
        "        batch_size, seq_len, d_model = x.shape\n",
        "\n",
        "        # Project input to keys and values\n",
        "        keys = self.key_proj(x)  # [batch_size, seq_len, d_model]\n",
        "        values = self.value_proj(x)  # [batch_size, seq_len, d_model]\n",
        "\n",
        "        # Compute attention scores with memory\n",
        "        attention_scores = torch.matmul(keys, self.memory.T)  # [batch_size, seq_len, memory_size]\n",
        "        attention_weights = F.softmax(attention_scores, dim=-1)\n",
        "\n",
        "        # Retrieve from memory\n",
        "        retrieved_memory = torch.matmul(attention_weights, self.memory)  # [batch_size, seq_len, d_model]\n",
        "\n",
        "        # Update memory based on surprise\n",
        "        surprise = torch.norm(values - retrieved_memory, dim=-1, keepdim=True)  # [batch_size, seq_len, 1]\n",
        "        forgetting_weights = torch.sigmoid(self.forgetting_gate(values))  # [batch_size, seq_len, 1]\n",
        "\n",
        "        # Update memory (during inference only)\n",
        "        if not self.training:\n",
        "            # Reduce batch and seq dimensions to match memory size\n",
        "            avg_forgetting_weights = forgetting_weights.mean(dim=(0, 1))  # [1, d_model]\n",
        "            avg_values = values.mean(dim=(0, 1))  # [1, d_model]\n",
        "\n",
        "            # Expand or reshape to match memory shape\n",
        "            avg_forgetting_weights = avg_forgetting_weights.unsqueeze(0).expand(self.memory.size(0), -1)  # [memory_size, d_model]\n",
        "            avg_values = avg_values.unsqueeze(0).expand(self.memory.size(0), -1)  # [memory_size, d_model]\n",
        "\n",
        "            # Update memory\n",
        "            self.memory.data = avg_forgetting_weights * self.memory + (1 - avg_forgetting_weights) * avg_values\n",
        "\n",
        "        return retrieved_memory\n",
        "\n",
        "\n",
        "class TitansTransformerEncoderLayer(nn.Module):\n",
        "    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, memory_size=512):\n",
        "        super().__init__()\n",
        "        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)\n",
        "        self.titans_memory = TitansMemoryModule(d_model, memory_size)\n",
        "\n",
        "        # Feed-forward network\n",
        "        self.linear1 = nn.Linear(d_model, dim_feedforward)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.linear2 = nn.Linear(dim_feedforward, d_model)\n",
        "\n",
        "        # Layer normalization\n",
        "        self.norm1 = nn.LayerNorm(d_model)\n",
        "        self.norm2 = nn.LayerNorm(d_model)\n",
        "        self.norm3 = nn.LayerNorm(d_model)\n",
        "        self.dropout1 = nn.Dropout(dropout)\n",
        "        self.dropout2 = nn.Dropout(dropout)\n",
        "        self.dropout3 = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, src, src_mask=None, src_key_padding_mask=None):\n",
        "        # Self-attention\n",
        "        src2 = self.self_attn(src, src, src, attn_mask=src_mask,\n",
        "                             key_padding_mask=src_key_padding_mask)[0]\n",
        "        src = src + self.dropout1(src2)\n",
        "        src = self.norm1(src)\n",
        "\n",
        "        # Titans memory integration\n",
        "        memory_output = self.titans_memory(src)\n",
        "        src = src + self.dropout2(memory_output)\n",
        "        src = self.norm2(src)\n",
        "\n",
        "        # Feed-forward network\n",
        "        src2 = self.linear2(self.dropout(F.relu(self.linear1(src))))\n",
        "        src = src + self.dropout3(src2)\n",
        "        src = self.norm3(src)\n",
        "\n",
        "        return src\n",
        "\n",
        "class TitansTransformer(nn.Module):\n",
        "    def __init__(self, num_tokens, d_model=512, nhead=8, num_layers=6,\n",
        "                 dim_feedforward=2048, dropout=0.1, memory_size=512):\n",
        "        super().__init__()\n",
        "\n",
        "        self.embedding = nn.Embedding(num_tokens, d_model)\n",
        "        self.pos_encoder = PositionalEncoding(d_model, dropout)\n",
        "\n",
        "        # Create encoder layers with Titans memory\n",
        "        self.layers = nn.ModuleList([\n",
        "            TitansTransformerEncoderLayer(d_model, nhead, dim_feedforward,\n",
        "                                        dropout, memory_size)\n",
        "            for _ in range(num_layers)\n",
        "        ])\n",
        "\n",
        "        self.norm = nn.LayerNorm(d_model)\n",
        "        self.fc_out = nn.Linear(d_model, num_tokens)\n",
        "\n",
        "        self.d_model = d_model\n",
        "        self._reset_parameters()\n",
        "\n",
        "    def _reset_parameters(self):\n",
        "        for p in self.parameters():\n",
        "            if p.dim() > 1:\n",
        "                nn.init.xavier_uniform_(p)\n",
        "\n",
        "    def forward(self, src, src_mask=None, src_key_padding_mask=None):\n",
        "        src = self.embedding(src) * math.sqrt(self.d_model)\n",
        "        src = self.pos_encoder(src)\n",
        "\n",
        "        for layer in self.layers:\n",
        "            src = layer(src, src_mask, src_key_padding_mask)\n",
        "\n",
        "        src = self.norm(src)\n",
        "        output = self.fc_out(src)\n",
        "        return output\n",
        "\n",
        "class PositionalEncoding(nn.Module):\n",
        "    def __init__(self, d_model, dropout=0.1, max_len=5000):\n",
        "        super().__init__()\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "\n",
        "        position = torch.arange(max_len).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n",
        "        pe = torch.zeros(max_len, 1, d_model)\n",
        "        pe[:, 0, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 0, 1::2] = torch.cos(position * div_term)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.pe[:x.size(0)]\n",
        "        return self.dropout(x)\n",
        "\n",
        "# Training utilities\n",
        "def create_mask(size):\n",
        "    mask = torch.triu(torch.ones(size, size) * float('-inf'), diagonal=1)\n",
        "    return mask\n",
        "\n",
        "def train_epoch(model, dataloader, optimizer, criterion, device):\n",
        "    model.train()\n",
        "    total_loss = 0\n",
        "\n",
        "    for batch in dataloader:\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        src = batch[:-1].to(device)\n",
        "        tgt = batch[1:].to(device)\n",
        "\n",
        "        mask = create_mask(src.size(1)).to(device)\n",
        "\n",
        "        output = model(src, src_mask=mask)\n",
        "        loss = criterion(output.view(-1, output.size(-1)), tgt.view(-1))\n",
        "\n",
        "        loss.backward()\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
        "        optimizer.step()\n",
        "\n",
        "        total_loss += loss.item()\n",
        "\n",
        "    return total_loss / len(dataloader)\n",
        "\n",
        "# Example usage:\n",
        "def main():\n",
        "    # Model parameters\n",
        "    num_tokens = 50000  # Vocabulary size\n",
        "    d_model = 512\n",
        "    nhead = 8\n",
        "    num_layers = 6\n",
        "    memory_size = 512\n",
        "\n",
        "    # Initialize model\n",
        "    model = TitansTransformer(\n",
        "        num_tokens=num_tokens,\n",
        "        d_model=d_model,\n",
        "        nhead=nhead,\n",
        "        num_layers=num_layers,\n",
        "        memory_size=memory_size\n",
        "    )\n",
        "\n",
        "    # Training setup\n",
        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "    model = model.to(device)\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "    print(\"Model initialized and ready for training\")\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ],
      "metadata": {
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mzHHrJhOn7xX",
        "outputId": "5774f2d1-07a1-4869-da00-200e16a478eb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model initialized and ready for training\n"
          ]
        }
      ],
      "execution_count": 8
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install datasets"
      ],
      "metadata": {
        "collapsed": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uZZFgs4tpKwh",
        "outputId": "8da15225-5659-43f7-b02d-1e1f1d2ca98d"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: datasets in /usr/local/lib/python3.11/dist-packages (3.2.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from datasets) (3.16.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.11/dist-packages (from datasets) (1.26.4)\n",
            "Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.11/dist-packages (from datasets) (17.0.0)\n",
            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.11/dist-packages (from datasets) (0.3.8)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.11/dist-packages (from datasets) (2.2.2)\n",
            "Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.11/dist-packages (from datasets) (2.32.3)\n",
            "Requirement already satisfied: tqdm>=4.66.3 in /usr/local/lib/python3.11/dist-packages (from datasets) (4.67.1)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.11/dist-packages (from datasets) (3.5.0)\n",
            "Requirement already satisfied: multiprocess<0.70.17 in /usr/local/lib/python3.11/dist-packages (from datasets) (0.70.16)\n",
            "Requirement already satisfied: fsspec<=2024.9.0,>=2023.1.0 in /usr/local/lib/python3.11/dist-packages (from fsspec[http]<=2024.9.0,>=2023.1.0->datasets) (2024.9.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.11/dist-packages (from datasets) (3.11.11)\n",
            "Requirement already satisfied: huggingface-hub>=0.23.0 in /usr/local/lib/python3.11/dist-packages (from datasets) (0.27.1)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.11/dist-packages (from datasets) (24.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.11/dist-packages (from datasets) (6.0.2)\n",
            "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (2.4.4)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (1.3.2)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (24.3.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (1.5.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (6.1.0)\n",
            "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (0.2.1)\n",
            "Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (1.18.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub>=0.23.0->datasets) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests>=2.32.2->datasets) (3.4.1)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests>=2.32.2->datasets) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests>=2.32.2->datasets) (2.3.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests>=2.32.2->datasets) (2024.12.14)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.11/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.11/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.11/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.17.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch.utils.data import DataLoader\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoTokenizer\n",
        "import time\n",
        "import torch.nn as nn\n",
        "\n",
        "\n",
        "def get_data(subset_size=1000):\n",
        "    # Load dataset from Hugging Face\n",
        "    dataset = load_dataset(\"wikitext\", \"wikitext-2-raw-v1\")\n",
        "    tokenizer = AutoTokenizer.from_pretrained('gpt2')\n",
        "\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "    def tokenize_function(examples):\n",
        "        return tokenizer(examples['text'], truncation=True, padding='max_length', max_length=512)\n",
        "\n",
        "    # Tokenize the dataset and limit the size\n",
        "    tokenized_dataset = dataset.map(\n",
        "        tokenize_function,\n",
        "        batched=True,\n",
        "        remove_columns=dataset['train'].column_names\n",
        "    )\n",
        "\n",
        "    # Select a subset of the dataset\n",
        "    train_data = tokenized_dataset['train'].select(range(subset_size))\n",
        "    val_data = tokenized_dataset['validation'].select(range(subset_size // 10))  # Smaller validation set\n",
        "    test_data = tokenized_dataset['test'].select(range(subset_size // 10))       # Smaller test set\n",
        "\n",
        "    # Convert to PyTorch tensors\n",
        "    train_data = torch.tensor(train_data['input_ids'], dtype=torch.long)\n",
        "    val_data = torch.tensor(val_data['input_ids'], dtype=torch.long)\n",
        "    test_data = torch.tensor(test_data['input_ids'], dtype=torch.long)\n",
        "\n",
        "    vocab_size = tokenizer.vocab_size\n",
        "\n",
        "    return train_data, val_data, test_data, vocab_size\n",
        "\n",
        "\n",
        "def batchify(data, batch_size, device):\n",
        "    # Divide data into batch_size parts\n",
        "    nbatch = data.size(0) // batch_size\n",
        "    data = data.narrow(0, 0, nbatch * batch_size)\n",
        "    data = data.view(batch_size, -1).t().contiguous()\n",
        "    return data.to(device)\n",
        "\n",
        "\n",
        "def evaluate(model, data_source, criterion, batch_size, device):\n",
        "    model.eval()\n",
        "    total_loss = 0.\n",
        "    with torch.no_grad():\n",
        "        for i in range(0, data_source.size(0) - 1, batch_size):\n",
        "            data = data_source[i:i + batch_size].to(device)\n",
        "            targets = data_source[i + 1:i + 1 + batch_size].to(device)\n",
        "            # Ensure input and target sizes match\n",
        "            if data.size(0) != targets.size(0):\n",
        "                break  # Skip incomplete batch\n",
        "            output = model(data)\n",
        "            total_loss += criterion(output.view(-1, output.size(-1)), targets.view(-1)).item()\n",
        "    return total_loss / (data_source.size(0) - 1)\n",
        "\n",
        "\n",
        "def train_model():\n",
        "    # Hyperparameters\n",
        "    batch_size = 16\n",
        "    eval_batch_size = 10\n",
        "    d_model = 512\n",
        "    nhead = 8\n",
        "    num_layers = 6\n",
        "    memory_size = 512\n",
        "    epochs = 3  # Fewer epochs for testing\n",
        "    subset_size = 1000  # Limit dataset size for testing\n",
        "\n",
        "    # Setup device\n",
        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "    # Get data (subset for testing)\n",
        "    train_data, val_data, test_data, vocab_size = get_data(subset_size=subset_size)\n",
        "\n",
        "    # Initialize model\n",
        "    model = TitansTransformer(\n",
        "        num_tokens=vocab_size,\n",
        "        d_model=d_model,\n",
        "        nhead=nhead,\n",
        "        num_layers=num_layers,\n",
        "        memory_size=memory_size\n",
        "    ).to(device)\n",
        "\n",
        "    # Batchify data\n",
        "    train_data = batchify(train_data, batch_size, device)\n",
        "    val_data = batchify(val_data, eval_batch_size, device)\n",
        "    test_data = batchify(test_data, eval_batch_size, device)\n",
        "\n",
        "    # Training setup\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n",
        "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)\n",
        "\n",
        "    # Training loop\n",
        "    best_val_loss = float('inf')\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        epoch_start_time = time.time()\n",
        "\n",
        "        # Train\n",
        "        model.train()\n",
        "        total_loss = 0.\n",
        "        for batch, i in enumerate(range(0, train_data.size(0) - 1, batch_size)):\n",
        "            data = train_data[i:i + batch_size].to(device)\n",
        "            targets = train_data[i + 1:i + 1 + batch_size].to(device)\n",
        "\n",
        "            # Ensure input and target sizes match\n",
        "            if data.size(0) != targets.size(0):\n",
        "                break  # Skip incomplete batch\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            output = model(data)\n",
        "            loss = criterion(output.view(-1, output.size(-1)), targets.view(-1))\n",
        "            loss.backward()\n",
        "            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
        "            optimizer.step()\n",
        "\n",
        "            total_loss += loss.item()\n",
        "\n",
        "            if batch % 100 == 0:\n",
        "                curr_loss = total_loss / (batch + 1)\n",
        "                print(f'| epoch {epoch + 1:3d} | batch {batch:3d} | '\n",
        "                      f'loss {curr_loss:5.2f}')\n",
        "\n",
        "\n",
        "        # Evaluate\n",
        "        val_loss = evaluate(model, val_data, criterion, eval_batch_size, device)\n",
        "        print('-' * 89)\n",
        "        print(f'| end of epoch {epoch + 1:3d} | time: {time.time() - epoch_start_time:5.2f}s | '\n",
        "              f'valid loss {val_loss:5.2f}')\n",
        "        print('-' * 89)\n",
        "\n",
        "        if val_loss < best_val_loss:\n",
        "            best_val_loss = val_loss\n",
        "            torch.save(model.state_dict(), 'titans_transformer_model.pt')\n",
        "\n",
        "        scheduler.step()\n",
        "\n",
        "    # Test\n",
        "    model.load_state_dict(torch.load('titans_transformer_model.pt'))\n",
        "    test_loss = evaluate(model, test_data, criterion, eval_batch_size, device)\n",
        "    print('=' * 89)\n",
        "    print(f'| End of training | test loss {test_loss:5.2f}')\n",
        "    print('=' * 89)\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    train_model()\n"
      ],
      "metadata": {
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yPV16wVZn7xY",
        "outputId": "a045fe7e-00f7-440b-e27d-cd590bba0f92"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "| epoch   1 | batch   0 | loss 11.05\n",
            "| epoch   1 | batch 100 | loss  4.72\n",
            "| epoch   1 | batch 200 | loss  3.27\n",
            "| epoch   1 | batch 300 | loss  2.55\n",
            "| epoch   1 | batch 400 | loss  2.26\n",
            "| epoch   1 | batch 500 | loss  2.05\n",
            "| epoch   1 | batch 600 | loss  1.89\n",
            "| epoch   1 | batch 700 | loss  1.81\n",
            "| epoch   1 | batch 800 | loss  1.75\n",
            "| epoch   1 | batch 900 | loss  1.76\n",
            "| epoch   1 | batch 1000 | loss  1.79\n",
            "| epoch   1 | batch 1100 | loss  1.72\n",
            "| epoch   1 | batch 1200 | loss  1.72\n",
            "| epoch   1 | batch 1300 | loss  1.64\n",
            "| epoch   1 | batch 1400 | loss  1.62\n",
            "| epoch   1 | batch 1500 | loss  1.58\n",
            "| epoch   1 | batch 1600 | loss  1.57\n",
            "| epoch   1 | batch 1700 | loss  1.55\n",
            "| epoch   1 | batch 1800 | loss  1.53\n",
            "| epoch   1 | batch 1900 | loss  1.53\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   1 | time: 108.01s | valid loss  0.14\n",
            "-----------------------------------------------------------------------------------------\n",
            "| epoch   2 | batch   0 | loss  3.25\n",
            "| epoch   2 | batch 100 | loss  1.53\n",
            "| epoch   2 | batch 200 | loss  1.62\n",
            "| epoch   2 | batch 300 | loss  1.41\n",
            "| epoch   2 | batch 400 | loss  1.37\n",
            "| epoch   2 | batch 500 | loss  1.32\n",
            "| epoch   2 | batch 600 | loss  1.27\n",
            "| epoch   2 | batch 700 | loss  1.27\n",
            "| epoch   2 | batch 800 | loss  1.27\n",
            "| epoch   2 | batch 900 | loss  1.32\n",
            "| epoch   2 | batch 1000 | loss  1.38\n",
            "| epoch   2 | batch 1100 | loss  1.35\n",
            "| epoch   2 | batch 1200 | loss  1.38\n",
            "| epoch   2 | batch 1300 | loss  1.32\n",
            "| epoch   2 | batch 1400 | loss  1.32\n",
            "| epoch   2 | batch 1500 | loss  1.29\n",
            "| epoch   2 | batch 1600 | loss  1.30\n",
            "| epoch   2 | batch 1700 | loss  1.30\n",
            "| epoch   2 | batch 1800 | loss  1.28\n",
            "| epoch   2 | batch 1900 | loss  1.30\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   2 | time: 108.31s | valid loss  0.15\n",
            "-----------------------------------------------------------------------------------------\n",
            "| epoch   3 | batch   0 | loss  3.29\n",
            "| epoch   3 | batch 100 | loss  1.50\n",
            "| epoch   3 | batch 200 | loss  1.60\n",
            "| epoch   3 | batch 300 | loss  1.40\n",
            "| epoch   3 | batch 400 | loss  1.36\n",
            "| epoch   3 | batch 500 | loss  1.31\n",
            "| epoch   3 | batch 600 | loss  1.26\n",
            "| epoch   3 | batch 700 | loss  1.26\n",
            "| epoch   3 | batch 800 | loss  1.26\n",
            "| epoch   3 | batch 900 | loss  1.31\n",
            "| epoch   3 | batch 1000 | loss  1.37\n",
            "| epoch   3 | batch 1100 | loss  1.34\n",
            "| epoch   3 | batch 1200 | loss  1.37\n",
            "| epoch   3 | batch 1300 | loss  1.31\n",
            "| epoch   3 | batch 1400 | loss  1.31\n",
            "| epoch   3 | batch 1500 | loss  1.28\n",
            "| epoch   3 | batch 1600 | loss  1.29\n",
            "| epoch   3 | batch 1700 | loss  1.29\n",
            "| epoch   3 | batch 1800 | loss  1.27\n",
            "| epoch   3 | batch 1900 | loss  1.29\n",
            "-----------------------------------------------------------------------------------------\n",
            "| end of epoch   3 | time: 108.19s | valid loss  0.15\n",
            "-----------------------------------------------------------------------------------------\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-10-e7ed4d7074d6>:146: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
            "  model.load_state_dict(torch.load('titans_transformer_model.pt'))\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "=========================================================================================\n",
            "| End of training | test loss  0.12\n",
            "=========================================================================================\n"
          ]
        }
      ],
      "execution_count": 10
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import time\n",
        "import matplotlib.pyplot as plt\n",
        "from torch.utils.data import DataLoader\n",
        "import math\n",
        "import numpy as np\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoTokenizer\n",
        "\n",
        "\n",
        "# Positional Encoding for Standard Transformer\n",
        "class PositionalEncoding(nn.Module):\n",
        "    def __init__(self, d_model, dropout=0.1, max_len=5000):\n",
        "        super().__init__()\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "\n",
        "        position = torch.arange(max_len).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n",
        "        pe = torch.zeros(max_len, 1, d_model)\n",
        "        pe[:, 0, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 0, 1::2] = torch.cos(position * div_term)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # The line below is changed to slice along dimension 0 up to x.size(0)\n",
        "        x = x + self.pe[:x.size(0), :, :]  # Adjust slicing to match input sequence length\n",
        "        return self.dropout(x)\n",
        "\n",
        "\n",
        "# Define Standard Transformer\n",
        "class StandardTransformer(nn.Module):\n",
        "    def __init__(self, num_tokens, d_model=512, nhead=8, num_layers=6,\n",
        "                 dim_feedforward=2048, dropout=0.1):\n",
        "        super().__init__()\n",
        "        self.d_model = d_model\n",
        "        self.embedding = nn.Embedding(num_tokens, d_model)\n",
        "        self.pos_encoder = PositionalEncoding(d_model, dropout)\n",
        "        self.transformer = nn.Transformer(\n",
        "            d_model=d_model,\n",
        "            nhead=nhead,\n",
        "            num_encoder_layers=num_layers,\n",
        "            num_decoder_layers=0,\n",
        "            dim_feedforward=dim_feedforward,\n",
        "            dropout=dropout,\n",
        "            batch_first=True\n",
        "        )\n",
        "        self.fc_out = nn.Linear(d_model, num_tokens)\n",
        "\n",
        "    def forward(self, src, src_mask=None):\n",
        "        src = self.embedding(src) * math.sqrt(self.d_model)\n",
        "        src = self.pos_encoder(src)\n",
        "        output = self.transformer.encoder(src, src_mask)\n",
        "        return self.fc_out(output)\n",
        "\n",
        "\n",
        "# Function to load data\n",
        "def get_data():\n",
        "    dataset = load_dataset(\"wikitext\", \"wikitext-2-raw-v1\")\n",
        "    tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
        "    tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "    def tokenize_function(examples):\n",
        "        return tokenizer(examples[\"text\"], truncation=True, padding=\"max_length\", max_length=512)\n",
        "\n",
        "    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=[\"text\"])\n",
        "\n",
        "    train_data = torch.tensor(tokenized_dataset[\"train\"][\"input_ids\"], dtype=torch.long)\n",
        "    val_data = torch.tensor(tokenized_dataset[\"validation\"][\"input_ids\"], dtype=torch.long)\n",
        "    test_data = torch.tensor(tokenized_dataset[\"test\"][\"input_ids\"], dtype=torch.long)\n",
        "\n",
        "    vocab_size = tokenizer.vocab_size\n",
        "    return train_data, val_data, test_data, vocab_size\n",
        "\n",
        "\n",
        "# Benchmarking function\n",
        "def benchmark_models(sequence_lengths=[128, 256, 512, 1024], batch_size=8):\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    results = {\n",
        "        \"standard\": {\"time\": [], \"memory\": [], \"perplexity\": []},\n",
        "        \"titans\": {\"time\": [], \"memory\": [], \"perplexity\": []},\n",
        "    }\n",
        "\n",
        "    train_data, val_data, test_data, vocab_size = get_data()\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "    for seq_len in sequence_lengths:\n",
        "        print(f\"\\nTesting sequence length: {seq_len}\")\n",
        "\n",
        "        # Initialize models\n",
        "        standard_transformer = StandardTransformer(num_tokens=vocab_size).to(device)\n",
        "        titans_transformer = TitansTransformer(num_tokens=vocab_size, memory_size=512).to(device)\n",
        "\n",
        "        # Prepare test batch\n",
        "        test_batch = test_data[:batch_size].to(device)\n",
        "\n",
        "        # Test Standard Transformer\n",
        "        torch.cuda.empty_cache()\n",
        "        start_time = time.time()\n",
        "        with torch.no_grad():\n",
        "            standard_transformer.eval()\n",
        "            output_standard = standard_transformer(test_batch)\n",
        "            perplexity_standard = torch.exp(\n",
        "                criterion(output_standard.view(-1, output_standard.size(-1)), test_batch.view(-1))\n",
        "            )\n",
        "        end_time = time.time()\n",
        "\n",
        "        results[\"standard\"][\"time\"].append(end_time - start_time)\n",
        "        results[\"standard\"][\"perplexity\"].append(perplexity_standard.item())\n",
        "\n",
        "        # Test Titans Transformer\n",
        "        torch.cuda.empty_cache()\n",
        "        start_time = time.time()\n",
        "        with torch.no_grad():\n",
        "            titans_transformer.eval()\n",
        "            output_titans = titans_transformer(test_batch)\n",
        "            perplexity_titans = torch.exp(\n",
        "                criterion(output_titans.view(-1, output_titans.size(-1)), test_batch.view(-1))\n",
        "            )\n",
        "        end_time = time.time()\n",
        "\n",
        "        results[\"titans\"][\"time\"].append(end_time - start_time)\n",
        "        results[\"titans\"][\"perplexity\"].append(perplexity_titans.item())\n",
        "\n",
        "        print(f\"Standard Transformer - Time: {results['standard']['time'][-1]:.4f}s, \"\n",
        "              f\"Perplexity: {results['standard']['perplexity'][-1]:.2f}\")\n",
        "        print(f\"Titans Transformer - Time: {results['titans']['time'][-1]:.4f}s, \"\n",
        "              f\"Perplexity: {results['titans']['perplexity'][-1]:.2f}\")\n",
        "\n",
        "    return results, sequence_lengths\n",
        "\n",
        "\n",
        "# Plot benchmark results\n",
        "def plot_benchmark_results(results, sequence_lengths):\n",
        "    fig, axes = plt.subplots(1, 2, figsize=(12, 5))\n",
        "\n",
        "    # Inference time\n",
        "    axes[0].plot(sequence_lengths, results[\"standard\"][\"time\"], label=\"Standard Transformer\", color=\"blue\")\n",
        "    axes[0].plot(sequence_lengths, results[\"titans\"][\"time\"], label=\"Titans Transformer\", color=\"red\")\n",
        "    axes[0].set_title(\"Inference Time\")\n",
        "    axes[0].set_xlabel(\"Sequence Length\")\n",
        "    axes[0].set_ylabel(\"Time (s)\")\n",
        "    axes[0].legend()\n",
        "\n",
        "    # Perplexity\n",
        "    axes[1].plot(sequence_lengths, results[\"standard\"][\"perplexity\"], label=\"Standard Transformer\", color=\"blue\")\n",
        "    axes[1].plot(sequence_lengths, results[\"titans\"][\"perplexity\"], label=\"Titans Transformer\", color=\"red\")\n",
        "    axes[1].set_title(\"Perplexity\")\n",
        "    axes[1].set_xlabel(\"Sequence Length\")\n",
        "    axes[1].set_ylabel(\"Perplexity\")\n",
        "    axes[1].legend()\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.savefig(\"benchmark_results.png\")\n",
        "    plt.close()\n",
        "\n",
        "\n",
        "# Main function\n",
        "if __name__ == \"__main__\":\n",
        "    print(\"\\nRunning benchmarks...\")\n",
        "    results, sequence_lengths = benchmark_models()\n",
        "    plot_benchmark_results(results, sequence_lengths)\n",
        "    print(\"\\nBenchmark results saved to 'benchmark_results.png'.\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZS2-2cF5oa-P",
        "outputId": "096cff57-99ec-4151-b03a-7e6e9a8b2b53"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Running benchmarks...\n",
            "\n",
            "Testing sequence length: 128\n",
            "Standard Transformer - Time: 0.1598s, Perplexity: 74703.84\n",
            "Titans Transformer - Time: 0.0329s, Perplexity: 63075.21\n",
            "\n",
            "Testing sequence length: 256\n",
            "Standard Transformer - Time: 0.0040s, Perplexity: 63956.66\n",
            "Titans Transformer - Time: 0.0078s, Perplexity: 56063.02\n",
            "\n",
            "Testing sequence length: 512\n",
            "Standard Transformer - Time: 0.0044s, Perplexity: 40087.42\n",
            "Titans Transformer - Time: 0.0086s, Perplexity: 40837.51\n",
            "\n",
            "Testing sequence length: 1024\n",
            "Standard Transformer - Time: 0.0038s, Perplexity: 31984.99\n",
            "Titans Transformer - Time: 0.0086s, Perplexity: 47995.84\n",
            "\n",
            "Benchmark results saved to 'benchmark_results.png'.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "M8ufG7_YudFN"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}